# ============================================================
# 该脚本用于自动识别衣服图片中的正反面，并将正反面配对拼接保存（按文件名前缀分组配对）
# ============================================================

f'''
逻辑更改：
# 伪代码描述：
# 1. 遍历所有图片文件名，只处理以.jpg结尾的文件。
# 2. 对每个文件名，去除扩展名后按下划线分割，判断最后一段是否为数字（编号）。
#    - 如果倒数第二段也是数字，则分组键为去掉最后两段的前缀，否则为去掉最后一段的前缀。
#    - 将同一分组键的图片路径收集到一起（即同一衣服的不同视角）。
# 3. 对每个分组（同一衣服的图片集合）：
#    - 如果该组图片数量小于2，跳过。
#    - 对组内每张图片，利用大模型判断其为正面还是反面，分别放入front和back列表。
#    - 如果正反面都存在，则两两配对，拼接图片（正面在左，反面在右）。
#    - 对拼接后的图片再次用大模型判断是否为同一件衣服的正反面，且左右颜色一致，才保存拼接图。
#    - 一旦找到一组满足条件的正反面配对即保存并跳出循环。

for n in names:
    if not n endswith .jpg:
        continue
    n_noext, ext = os.path.splitext(n)
    n_split = n_noext.split('_')
    if n_split[-1] is digital:
        if n_split[-2] is digital:
            key = '_'.join( n_split[:-2] )
        else:
            key = '_'.join( n_split[:-1] )
        dict[key] .append osj( data_dirk , n ) # 将文件名前缀相同的 路径放在一起

for k,v in dict:
    if v len < 2:continue
    front = []
    back = []
    for path in v:
        ans_front from path
        ans_back from path
        if check front:
            front append (path,pil)
        elif check back:
            fron append (path,pil)
        else:pass

        if front and back: # back 数量 一般比 front 少
            for b in back:
                for f in front:
                    f_path, f_pil = f
                    b_path, b_pil = b

                    f_name = osb(f_path)
                    b_name = osb(b_path)
                    
                    concat_img = concat(f_pil, b_pil)
                    save_name = f'f_name__b_name.jpg'
                    save_path = osj( save_dir, save_name )
                    if check ans is front_back:
                        if check ans is same color:
                            concat_img.save( save_path )

                            break
'''

import os
import shutil
from PIL import Image
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='1', help='CUDA device id(s) to use')
args, _ = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

from pathlib import Path
import cv2
import numpy as np
from utils.util_for_os import osj, ose, osb

# 获取数据和保存目录的lambda函数
get_data_dir = lambda category, date: f'/mnt/nas/datasets/diction/{category}{date}_img_clo'  # sifted clo
get_save_dir = lambda category, date: f'/mnt/nas/datasets/diction/{category}{date}_img_front_back'  # img

# 只处理sweater类别和指定日期
# categories = ['coat']
# dates = ['0903','0904',
#         '0909','0910']
# categories = ['sweater']
# dates = ['0828','0829',
#         '0901','0902','0903','0904','0905',
#         '0908','0909','0910']

# categories = ['dress', 'slipdress', 'feather']
dates = ['1106']
categories = ['jacket','hoodie','formalatt','weddress']
# dates = ['0918', '0919']
# categories = ['dress','slipdress']
# dates = ['0925','0926','0928']

# 检查回答是否为否/是的辅助函数
check_no = lambda ans: 'no' in ans.lower() or '否' in ans or '不是' in ans
check_yes = lambda ans: 'yes' in ans.lower() or '是' in ans

import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import VLChatProcessor

# 加载Janus多模态大模型和处理器
model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config
language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(
    model_path,
    language_config=language_config,
    trust_remote_code=True
)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'

@torch.inference_mode()
def multimodal_understanding(image, question, seed=42, top_p=0.95, temperature=0.1):
    """
    使用Janus模型对图片和问题进行多模态理解，返回模型回答
    """
    # 清理CUDA缓存
    torch.cuda.empty_cache()

    # 设置随机种子，保证结果可复现
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)

    # 构造对话输入
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]

    pil_images = [Image.fromarray(image)]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)

    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)

    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )

    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

def get_answer(filepath, q, img_pil=None):
    """
    获取图片的模型回答
    - filepath: 图片路径
    - q: 问题
    - img_pil: 可选，已加载的PIL图片
    """
    if img_pil is not None:
        np_image = np.array(img_pil)  # shape: (H, W, 3)
        # RGB转BGR（OpenCV默认格式）
        cv2_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
    else:
        cv2_image = cv2.imread(filepath)
        if cv2_image is None:
            # 错误处理：图片路径无效或图片损坏
            print(f"Error: Cannot read image {filepath}")
            return ""
    return multimodal_understanding(cv2_image, q)

# 问题定义
Q8 = 'Is a front view from cloth'
Q9 = 'Is a back view from cloth'
def is_clo_front(ans_front, ans_back):
    """
    判断是否为衣服正面
    - ans_front: 回答Q8
    - ans_back: 回答Q9
    """
    if check_yes(ans_front) and check_no(ans_back):
        return True
    return False

def is_clo_back(ans_front, ans_back):
    """
    判断是否为衣服背面
    - ans_front: 回答Q8
    - ans_back: 回答Q9
    """
    if check_no(ans_front) and check_yes(ans_back):
        return True
    return False

Q10 = 'Is a front view and a back view from the same cloth'
Q11 = "What's the color of the left clo? What's the color of the right clo? Is the same color between the left and the right ?"
def check_Q10(ans):
    """
    判断拼接图像是否为同一件衣服的正反面
    """
    if check_yes(ans):
        return True
    return False

def check_Q11(ans):
    """
    判断左右衣服颜色是否相同
    """
    if 'same' in ans.lower() and 'the same color is not' not in ans.lower():
        return True
    return False

from tqdm import tqdm
from utils.util_flux import process_img_1024, horizontal_concat_images

# 1. 检查数据目录和names.txt文件是否存在
for category in categories:
    for date in dates:
        if date >= '0826' and category == 'leather':
            continue
        data_dir = get_data_dir(category, date)
        name_path = osj(data_dir, 'names.txt')
        assert ose(name_path), f"names.txt not found: {name_path}"

def is_digit(s):
    """
    判断字符串s是否为数字
    """
    try:
        int(s)
        return True
    except:
        return False

for category in categories:
    for date in dates:
        if date >= '0826' and category == 'leather':
            continue

        data_dir = get_data_dir(category, date)
        name_path = osj(data_dir, 'names.txt')
        assert ose(name_path), f"names.txt not found: {name_path}"
        with open(name_path, encoding='utf-8') as f:
            jpg_names = f.readlines()
        # 只保留.jpg结尾的文件名
        jpgnames = [name.strip() for name in jpg_names if name.strip().endswith('.jpg')]

        save_dir = get_save_dir(category, date)
        if ose(save_dir):
            shutil.rmtree(save_dir)
        os.makedirs(save_dir)

        # 1. 按文件名前缀分组
        from collections import defaultdict
        group_dict = defaultdict(list)
        for n in tqdm(jpgnames, desc='prefix'):
            # 分割文件名，去除扩展名
            n_noext, ext = os.path.splitext(n)
            n_split = n_noext.split('_')
            if len(n_split) < 2:
                continue
            # 判断文件名最后一段是否为数字（编号）
            if is_digit(n_split[-1]):
                # 如果倒数第二段也是数字，则前缀去掉最后一段，否则保留全部
                if len(n_split) >= 2 and is_digit(n_split[-2]):
                    key = '_'.join(n_split[:-1])
                else:
                    key = '_'.join(n_split)
                group_dict[key].append(osj(data_dir, n))
        
        
        # # 可选：打印分组信息用于调试
        # print("==== group_dict 内容预览 ====")
        # for key, paths in group_dict.items():
        #     if len(paths) < 2:continue
        #     print(f"组名: {key}")
        #     print(f"  文件数: {len(paths)}")
        #     for p in paths:
        #         print(f"    {p}")
        #     print("-" * 40)
        # print("==== group_dict 打印结束 ====")
        
        # exit(1)


        # 2. 对每组进行正反面识别与配对
        for k, v in tqdm(group_dict.items(), desc="group"):
            if len(v) < 2:
                continue
            front = []
            back = []
            for path in v:
                # 利用大模型判断正反面
                ans_front = get_answer(path, Q8)
                ans_back = get_answer(path, Q9)
                pil = process_img_1024(path)
                if is_clo_front(ans_front, ans_back):
                    front.append((path, pil))
                elif is_clo_back(ans_front, ans_back):
                    back.append((path, pil))
                else:
                    pass
            # 配对正反面图片
            if front and back:
                for b in reversed(back):
                    for f in reversed(front):
                        f_path, f_pil = f
                        b_path, b_pil = b
                        f_name = osb(f_path)
                        b_name = osb(b_path)
                        # 拼接图片
                        concat_img = horizontal_concat_images([f_pil, b_pil])
                        save_name = f'{f_name}__{b_name}.jpg'
                        save_path = osj(save_dir, save_name)
                        # 判断拼接后是否为同一件衣服的正反面
                        ans = get_answer('', Q10, img_pil=concat_img)
                        if check_Q10(ans):
                            # 再问模型：左右颜色是否一致
                            ans_color = get_answer('', Q11, img_pil=concat_img)
                            if check_Q11(ans_color):
                                concat_img.save(save_path)
                                concat_img.save('tmp_concat2.jpg')
                                print('concat img find !!!!!', flush=True)
                                print(save_name, ans_color, flush=True)

                                # 确保 匹配后 这张 front 不再参与后续匹配
                                front.remove(f)
                                break

'''
date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"
bash restore.sh /mnt/nas/datasets/diction/suit"$date"
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo_diff

'''